
import numpy as np

from .panda import PandaPickAndPlaceEnv,PandaPushEnv

class MultiTaskPickAndPlaceEnv(PandaPickAndPlaceEnv):
    def __init__(self, task={}, n_tasks=2, **kwargs):
        self._task = task
        self.tasks = self.sample_tasks(n_tasks)
        self._goal = self.tasks[0]['goal']
        super(MultiTaskPickAndPlaceEnv, self).__init__(**kwargs)

    def get_all_task_idx(self):
        return range(len(self.tasks))

    def reset_task(self, idx):
        self._task = self.tasks[idx]
        self._goal = self._task['goal'] # assume parameterization of task by single vector
        self.reset_env(**self._goal)
        self.reset()

    def reset_env(self, lateral_friction=1.0,spinning_friction=0.001,mass=1.0,gravity=-9.81):
        block_uid = self.sim._bodies_idx['object']
        self.sim.physics_client.changeDynamics(bodyUniqueId=block_uid, linkIndex=-1, mass=mass)
        # change table's friction
        self.sim.set_lateral_friction('table', -1, lateral_friction=lateral_friction)
        self.sim.set_spinning_friction('table', -1, spinning_friction=spinning_friction)
        # gravity
        self.sim.physics_client.setGravity(0, 0, gravity)

class MultiTaskPushEnv(PandaPushEnv):

    def __init__(self, task={}, n_tasks=2, **kwargs):
        self._task = task
        self.tasks = self.sample_tasks(n_tasks)
        self._goal = self.tasks[0]['goal']
        super(MultiTaskPushEnv, self).__init__(**kwargs)

    def get_all_task_idx(self):
        return range(len(self.tasks))

    def reset_task(self, idx):
        self._task = self.tasks[idx]
        self._goal = self._task['goal'] # assume parameterization of task by single vector

        self.reset_env(**self._goal)
        self.reset()
    
    def reset_env(self, lateral_friction=1.0,spinning_friction=0.001,mass=1.0,gravity=-9.81):
        block_uid = self.sim._bodies_idx['object']
        self.sim.physics_client.changeDynamics(bodyUniqueId=block_uid, linkIndex=-1, mass=mass)
        # change table's friction
        self.sim.set_lateral_friction('table', -1, lateral_friction=lateral_friction)
        self.sim.set_spinning_friction('table', -1, spinning_friction=spinning_friction)
        # gravity
        self.sim.physics_client.setGravity(0, 0, gravity)